# %%
"""Functions to simulate CSTPU"""
# Not used, but interesting info:
#
# https://github.com/mvernacc/aerospike-nozzle-design-gui/blob/master/nozzle_solver.py
# mach_from_er
#     https://github.com/mvernacc/proptools/blob/master/proptools/nozzle.py
# Pe = params.Pc * (1 + (y-1)/2*Me**2)**(-y/(y-1))
# Документация - https://proptools.readthedocs.io/en/latest/nozzle_tutorial.html
#     Исходники документации с формулами в Latex
#         https://github.com/mvernacc/proptools/blob/master/docs/source/nozzle_tutorial.rst

from math import sqrt

import numpy as np
from pyXSteam.XSteam import XSteam
from scipy.optimize import Bounds, minimize

from mission_parameters import F_KR_0, FE_NOZZLE_0, T_INI_K, TANK_STRUCTURE_MASS, V_SUM

# TODO import constants from scipy or astropy:
# sigma, Stefan-Boltzmann constant 	5.670373e-08 W m^-2 K^-4


STEAM_TABLE = XSteam(XSteam.UNIT_SYSTEM_MKS)  # m/kg/sec/C/bar/W

# global constants definition
R = 461.889  # gas constant for H20, J/(K*moll)
GAMMA = 4.0 / 3.0  # adiabatic constant, explicitly float for Cython
C_ALUM = 903  # aluminum thermal capacity J/(kg*K)

VALVE_OPENING_TEMPERATURE = 450  # Kelvin
# To prevent thrusting start when temperature is too low, it's possible when
# satellite is waiting for thrusting allowed orbit zone
MIN_VALVE_OPENING_TEMPERATURE = VALVE_OPENING_TEMPERATURE - 10  # Kelvin

SUN_CONSTANT = 1367  # W/m^2

SIGMA = 5.67e-8  # Stephan-Boltzmann constant, W*m^2*K^4
EPSILON = 0.83  # for special space painting, https://bit.ly/2uAEKXl

PHI_C = 0.98
# the coefficient of the chamber,
# taking into account the loss in the chamber

PHI_N = 0.97
# The nozzle coefficient,
# taking into account the loss in the nozzle


# %%
def limit_to_range(num: float, range_min: float, range_max: float) -> float:
    """Limiting the number to the specified range.

    Parameters
    ----------
    num : float
        The number whose value should be limited.
    range_min : float
        The left border of the range.
    range_max : float
        The right border of the range.

    Returns
    -------
    float
        range_min if num<range_min, range_max if num>range_max, num in other cases.
    """
    if num < range_min:
        return range_min
    elif num > range_max:
        return range_max
    else:
        return num


# %%
def steam_enthalpy_mass_flow_product(temp_kelvin, mass_flow):
    """Calculates the energy loss of a water-steam mixture as
    equal to the thermal energy of released vapor.

    Parameters
    ----------
    temp_kelvin : float
        Saturated steam Temperature, Kelvin.
    mass_flow : float
        mass flow rate, kg/second.

    Returns
    -------
    res : float
        steam enthalpy and mdot product, Watt.
    """

    # 273.15 is water triple point temperature
    # 647.3 is water critical temperature
    # in this range both liquid water and steam exist
    temp_kelvin = limit_to_range(temp_kelvin, 273.5, 647)

    # XSteam library require temperature in Celsius!
    temp_celsius = temp_kelvin - 273.15  # Saturated steam Temperature, Celsius

    # TODO maybe need to change mass_flow > epsilon (10e-6 or another)
    if mass_flow > 0:
        res = mass_flow * STEAM_TABLE.hV_t(temp_celsius) * 1000
        # Saturated vapour enthalpy, J/kg
        # XSteam library calculates enthalpy in kJ/kg
    else:
        res = 0

    return res


# %%
def energy_err_squared(x_vector):
    """Calculates squared difference between actual internal energy
    and internal energy as function of supposed temperature.

    Parameters
    ----------
    x_vector : numpy array consists of
        temp_sup_kelvin : float
            Saturated steam Temperature, kelvin;
        m_sum : float
            Water-Steam mixture total mass, kg;
        internal_energy : float
            Water-Steam mixture thermal energy, J;

    Returns
    -------
    squared_err: float
    """
    temp_sup_kelvin, m_sum, internal_energy = x_vector

    temp_sup_celsius = temp_sup_kelvin - 273.15
    # Saturated steam Temperature, Celsius
    # XSteam library require temperature in Celsius

    rho_steam = STEAM_TABLE.rhoV_t(temp_sup_celsius)
    # Saturated steam density, kg/m**3
    rho_liq = STEAM_TABLE.rhoL_t(temp_sup_celsius)
    # Saturated liquid density, kg/m**3
    m_steam = rho_steam * (V_SUM * rho_liq - m_sum) / (rho_liq - rho_steam)
    # Steam mass, kg
    # because V_SUM = V_liq + V_steam
    # and m_sum = m_liq + m_steam

    m_liq = m_sum - m_steam
    # Water mass, kg
    h_steam = STEAM_TABLE.hV_t(temp_sup_celsius) * 1000
    # Saturated vapour enthalpy, J/kg
    h_liq = STEAM_TABLE.hL_t(temp_sup_celsius) * 1000
    # Saturated liquid enthalpy, J/kg
    # XSteam library calculates enthalpy in kJ/kg
    # TODO check previous statement and necessity for multiplication on 1000
    squared_err = (
        m_steam * h_steam
        + m_liq * h_liq
        + temp_sup_kelvin * C_ALUM * TANK_STRUCTURE_MASS
        - internal_energy
    ) ** 2

    return squared_err


# %%
def dry_steam_energy_err_squared(x_vector):
    """In a case when there is no liquid water remained in a tank,
    function calculates squared difference between actual internal energy
    and internal energy calculated as a function of supposed temperature.

    Parameters
    ----------
    x_vector : numpy array consists of
        temp_sup_kelvin : float
            steam Temperature, kelvin
        m_sum : float
            Water-Steam mixture mass, kg
        internal_energy : float
            Water-Steam mixture thermal energy, J

    Returns
    -------
    squared_err: float
    """

    temp_sup_kelvin, m_sum, internal_energy = x_vector

    temp_sup_celsius = temp_sup_kelvin - 273.15
    # steam Temperature, Celsius
    # XSteam library require temperature in Celsius

    m_steam = m_sum
    # Steam mass, kg

    h_steam = STEAM_TABLE.hV_t(temp_sup_celsius) * 1000
    # Saturated vapour enthalpy, J/kg
    # XSteam library calculates enthalpy in kJ/kg
    squared_err = (
        m_steam * h_steam
        + temp_sup_kelvin * C_ALUM * TANK_STRUCTURE_MASS
        - internal_energy
    ) ** 2

    return squared_err


# %%
def determine_valve_output(pressure_in, temperature_in, valve_opened):
    """Determines pressure and temperature behind valve based on valve state

    Parameters
    ----------
    pressure_in : float
        gas pressure above valve, Pa.
    temperature_in : float
        gas temperature above valve, Pa.
    valve_opened : boolean
        True if valve opened.

    Returns
    -------
    pressure_out : float
        gas pressure behind valve, Pa.
    temperature_out : float
        gas temperature behind valve, Pa.
    """

    if valve_opened:
        pressure_out = pressure_in
        temperature_out = temperature_in
    else:
        pressure_out = 0
        temperature_out = 0

    return pressure_out, temperature_out


# %%
def fe_nozzle_pa_err_squared(x_vector):
    """Calculates squared difference between actual nozzle exit area Fa and
    calculated nozzle exit area based on Pa and other gas parameters

    Parameters
    ----------
    x_vector : numpy array
        fe_nozzle : float
            nozzle exit area, m^2.
        mdot : float
            mass flow, m^2.
        pa_bar : float
            ambient pressure, bar.
        p_c_bar : float
            pressure in chamber, bar.
        temp_c_kelvin : float
            temperature in chamber, K.

    Returns
    -------
    squared_err : float
        squared and scaled error
    """
    fe_nozzle, mdot, pa_bar, p_c_bar, temp_c_kelvin = x_vector
    p_k = p_c_bar * 1e5
    p_a = pa_bar * 1e5
    # to reduce the difference in orders of magnitude
    # and reduce rounding errors

    squared_err = (
        (
            fe_nozzle
            - (mdot * (p_a / p_k) ** ((-1 + GAMMA) / GAMMA) * R * temp_c_kelvin)
            / (
                sqrt(2)
                * p_a
                * sqrt(
                    (
                        GAMMA
                        * R
                        * (
                            temp_c_kelvin
                            - (p_a / p_k) ** ((-1 + GAMMA) / GAMMA) * temp_c_kelvin
                        )
                    )
                    / (-1 + GAMMA)
                )
            )
        )
        * 1e8
    ) ** 2

    return squared_err


# %%
def calculate_flow_from_nozzle(p_c_bar, temp_c_kelvin):
    """Calculates flow parameters in a nozzle exit section
    based on gas parameters in chamber.

    Parameters
    ----------

    p_c_bar : float
        pressure in the chamber, bar.
    temp_c_kelvin : float
        temperature in the chamber, K.

    Returns
    -------

    thrust : float
        thrust, N.
    imp_sp : float
        specific impulse, m/sec.
    mass_flow : float
        mass flow, kg/sec.
    """

    p_k = p_c_bar * 1e5
    # to reduce the difference in orders of magnitude
    # and reduce rounding errors

    # ***********************************************
    # Setting the geometric parameters of the nozzle,
    # the constant parameters of the gas
    # and environmental parameters

    # pa_first_approx = 1e4
    # Thrust result became different after pa_first_approx changed
    pa_first_approx = 1e-4
    # initial approximation of pressure at the nozzle exit, Pa

    f_cr = F_KR_0 * PHI_N
    # The calculated area of the critical section,
    # taking into account losses, m**2

    fe_nozzle = FE_NOZZLE_0 * PHI_N**2
    # The calculated area of the output section of the nozzle,
    # taking into account losses, m**2

    # TODO maybe need to change p_k > epsilon (10e-6 or another)
    if p_k > 0:
        betta0 = sqrt(R * temp_c_kelvin) / (
            sqrt(GAMMA) * (2 / (GAMMA + 1)) ** ((GAMMA + 1) / (2 * (GAMMA - 1)))
        )
        # consumable parameter, m/s

        betta = betta0 * PHI_C
        # consumable parameter regarding losses, m/s

        mass_flow = p_k * f_cr / betta

        x_0 = np.array([fe_nozzle, mass_flow, pa_first_approx, p_k, temp_c_kelvin])
        l_b = np.array([fe_nozzle, mass_flow, 0, p_k, temp_c_kelvin])
        u_b = np.array([fe_nozzle, mass_flow, 3.19e-2, p_k, temp_c_kelvin])
        # u_b = np.array([fe_nozzle, mass_flow, 8.53e-5, p_k, temp_c_kelvin])
        # external pressure in range 0..3.19e-2 (equal to 100 km height)
        # external pressure in range 0..8.53e-5 (equal to 200 km height)
        # https://ru.wikipedia.org/wiki/%D0%A1%D1%82%D0%B0%D0%BD%D0%B4%D0%B0%D1%80%D1%82%D0%BD%D0%B0%D1%8F_%D0%B0%D1%82%D0%BC%D0%BE%D1%81%D1%84%D0%B5%D1%80%D0%B0

        bounds = Bounds(l_b, u_b)
        res = minimize(
            fe_nozzle_pa_err_squared,
            x_0,
            method="Powell",
            bounds=bounds,
            tol=1e-4,
        )
        p_a = res.x[2]

        w_cr = sqrt(2 * GAMMA / (GAMMA + 1) * R * temp_c_kelvin)
        # gas velocity in the critical section, m/s
        w_a = sqrt(
            2
            * GAMMA
            / (GAMMA - 1)
            * R
            * temp_c_kelvin
            * (1 - (p_a / p_k) ** ((GAMMA - 1) / GAMMA))
        )
        # outflow velocity, m/s
        lambda_a = w_a / w_cr  # speed ratio
        ktp = (2 / (GAMMA + 1)) ** (1 / (GAMMA - 1)) * (1 + lambda_a**2) / lambda_a
        # The coefficient of thrust in vacuum
        thrust = f_cr * ktp * p_k  # thrust, N
        epsilon = p_k / p_a
        phi1 = PHI_C * PHI_N
        # Coefficient taking into account losses in the chamber
        # and nozzle
        w_a_r = phi1 * w_a  # actual flow rate, m/s
        imp_sp = w_a_r + FE_NOZZLE_0 / F_KR_0 * betta / epsilon
        # specific void impulse, m/s
    else:
        thrust = 0
        imp_sp = 0
        mass_flow = 0

    return thrust, imp_sp, mass_flow


# %%
def calc_water_steam_equilibrium(internal_energy, m_sum):
    """Calculates temperature, pressure, steam and liquid water masses as a
    function of internal energy.

    Parameters
    ----------
    internal_energy : float
        Water-steam-tank total thermal energy, J.
    m_sum : float
        Water-steam total mass, kg.

    Returns
    -------
    p_bar : float
        pressure in chamber, bar.
    t_real_k : float
        real temperature in chamber, K.
    m_steam : float
        steam mass in chamber, kg.
    m_liq : float
        liquid water mass in chamber, kg.
    """

    x_0 = np.array([T_INI_K, m_sum, internal_energy])

    # 273.15 is water triple point temperature
    # 647.3 is water critical temperature
    # in this range both liquid water and steam exist
    u_b = np.array([min(VALVE_OPENING_TEMPERATURE + 1, 647), m_sum, internal_energy])
    l_b = np.array([273.15 + 1, m_sum, internal_energy])

    bounds = Bounds(l_b, u_b)
    res = minimize(
        energy_err_squared,
        x_0,
        method="Powell",
        bounds=bounds,
        tol=1e-4,
    )
    t_real_k = res.x[0]

    p_bar = STEAM_TABLE.psat_t(t_real_k - 273.15)  # Saturated steam pressure, bar

    rho_steam = STEAM_TABLE.rhoV_t(t_real_k - 273.15)
    # Saturated steam density, kg/m**3
    rho_liq = STEAM_TABLE.rhoL_t(t_real_k - 273.15)
    # Saturated liquid density, kg/m**3
    m_steam = rho_steam * (V_SUM * rho_liq - m_sum) / (rho_liq - rho_steam)
    # Steam mass, kg
    m_liq = m_sum - m_steam  # Water mass, kg
    if m_liq <= 0:
        # in this case only dry steam in tank without liquid water
        res = minimize(
            dry_steam_energy_err_squared,
            x_0,
            method="Powell",
            bounds=bounds,
            tol=1e-4,
        )
        t_real_k = res.x[0]

        p_pa = m_sum / V_SUM * R * t_real_k
        p_bar = p_pa * 1e-5
        m_steam = m_sum
        m_liq = 0

    return p_bar, t_real_k, m_steam, m_liq


# %%
def internal_energy_from_temperature(m_sum, temp_kelvin):
    """Returns water-steam mixture internal energy as a function of temperature.

    Parameters
    ----------
    m_sum : float
        Water-steam total mass, kg.
    temp_kelvin : float
        temperature, K.


    Returns
    -------
    internal_energy: float
        internal energy, J.
    """

    temp_celsius = temp_kelvin - 273.15

    rho_steam = STEAM_TABLE.rhoV_t(temp_celsius)
    # Saturated steam density, kg/m**3
    rho_liq = STEAM_TABLE.rhoL_t(temp_celsius)
    # Saturated liquid density, kg/m**3
    m_steam = rho_steam * (V_SUM * rho_liq - m_sum) / (rho_liq - rho_steam)
    # because V_SUM = V_liq + V_steam
    # and m_sum = m_liq + m_steam
    m_liq = m_sum - m_steam

    if m_liq > 0:
        h_steam = STEAM_TABLE.hV_t(temp_celsius) * 1000
        # Saturated steam enthalpy, J/kg
        h_liq = STEAM_TABLE.hL_t(temp_celsius) * 1000
        # Saturated liquid enthalpy, J/kg
        # XSteam library calculates enthalpy in kJ/kg
        internal_energy = (
            m_steam * h_steam
            + m_liq * h_liq
            + temp_kelvin * C_ALUM * TANK_STRUCTURE_MASS
        )
    else:
        p_pa = m_sum / V_SUM * R * temp_kelvin  # Pressure, Pa
        m_steam = m_sum
        h_steam = STEAM_TABLE.h_pt(p_pa * 1e-5, temp_celsius) * 1000
        # Enthalpy as a function of pressure(bar) and temperature(C), J/kg
        internal_energy = m_steam * h_steam + temp_kelvin * C_ALUM * TANK_STRUCTURE_MASS

    return internal_energy
